//
// Created on 2025/6/19.
//
// Node APIs are not fully supported. To solve the compilation error of the interface cannot be found,
// please include "napi/native_api.h".

#ifndef MSLITETOKENIZER_OPUS_MT_TOKENIZER_H
#define MSLITETOKENIZER_OPUS_MT_TOKENIZER_H

#include "common.h"
#include <sentencepiece/sentencepiece_processor.h>
#include <jsoncpp/json.h>
#include <fstream>
#include <vector>
#include <map>

#include <sentencepiece/sentencepiece_processor.h>
#include <jsoncpp/json.h>
#include <fstream>
#include <vector>
#include <string>
#include <map>
#include <algorithm>
#include <cctype>

class MarianTokenizer {
public:
    // 构造函数，加载所有必要资源
    MarianTokenizer() {
        // 加载SentencePiece模型
        source_sp_.Load(tokenizer_dir + "/source.spm");
        target_sp_.Load(tokenizer_dir + "/target.spm");
        
        // 加载tokenizer配置
        loadConfig(tokenizer_dir + "/tokenizer_config.json");
        
        // 加载词汇表
        loadVocab(tokenizer_dir + "/vocab.json");
        
        // 从vocab加载特殊token
        unk_token_id_ = getTokenId("<unk>");
        pad_token_id_ = getTokenId("<pad>");
        eos_token_id_ = getTokenId("</s>");
    }
    
    // 中文编码函数（源语言处理）
    std::vector<int> encode(const std::string& src_text) {
        // 添加语言标记：">>eng<<" (目标语言标记)
        // std::string full_text = ">>" + target_lang_ + "<< " + src_text;
        
        // 分词并转换为id
        std::vector<int> token_ids;
        std::vector<std::string> pieces;
        
        // 使用SentencePiece分词
        source_sp_.Encode(src_text, &pieces);
        
        // 将分词语块转换为id
        for (const auto& piece : pieces) {
            int id = getTokenId(piece);
            token_ids.push_back(id);
        }
        
        return token_ids;
    }
    
    // 英文解码函数（目标语言处理）
    std::string decode(const std::vector<int>& token_ids) {
        std::vector<std::string> pieces;
        
        for (int id : token_ids) {
            // 跳过特殊token
            if (id == eos_token_id_ || id == pad_token_id_) 
                continue;
            
            // 获取token对应的文本
            std::string token_str = getTokenString(id);
            
            // 特殊处理：target.spm使用▁表示空格
            if (token_str == "<unk>") {
                pieces.push_back("▁");  // 替换UNK
            } else {
                pieces.push_back(token_str);
            }
        }
        
        // 使用SentencePiece解码器
        std::string decoded;
        target_sp_.Decode(pieces, &decoded);
        
        // 后处理：删除多余空格和特殊字符
        return postProcess(decoded);
    }
    
private:
    // 加载tokenizer配置
    void loadConfig(const std::string& configPath) {
        std::ifstream configFile(configPath);
        Json::Value root;
        configFile >> root;
        
        source_lang_ = root["source_lang"].asString();
        target_lang_ = root["target_lang"].asString();
    }
    
    // 加载词汇表
    void loadVocab(const std::string& vocabPath) {
        std::ifstream vocabFile(vocabPath);
        Json::Value root;
        vocabFile >> root;
        
        for (auto it = root.begin(); it != root.end(); ++it) {
            std::string token = it.key().asString();
            int id = it->asInt();
            token_to_id_[token] = id;
            id_to_token_[id] = token;
        }
    }
    
    // 获取token ID
    int getTokenId(const std::string& token) {
        auto it = token_to_id_.find(token);
        if (it != token_to_id_.end()) {
            return it->second;
        }
        return unk_token_id_;  // 未知token
    }
    
    // 获取token文本
    std::string getTokenString(int id) {
        auto it = id_to_token_.find(id);
        if (it != id_to_token_.end()) {
            return it->second;
        }
        return "<unk>";  // 未知ID
    }
    
    // 后处理解码结果
    std::string postProcess(std::string text) {
        // 替换特殊空格表示
        size_t pos = 0;
        while ((pos = text.find("▁", pos)) != std::string::npos) {
            text.replace(pos, 3, " ");  // ▁是3字节UTF-8字符
        }
        
        // 删除多余空格
        auto end = std::unique(text.begin(), text.end(),
            [](char l, char r){ return std::isspace(l) && std::isspace(r); });
        text.erase(end, text.end());
        
        // 修剪首尾空格
        size_t start = text.find_first_not_of(" ");
        size_t end_pos = text.find_last_not_of(" ");
        if (start == std::string::npos || end_pos == std::string::npos)
            return "";
        return text.substr(start, end_pos - start + 1);
    }
    
    sentencepiece::SentencePieceProcessor source_sp_;  // 中文处理模型
    sentencepiece::SentencePieceProcessor target_sp_;  // 英文处理模型
    
    std::string source_lang_;
    std::string target_lang_;
    
    std::map<std::string, int> token_to_id_;  // token→ID映射
    std::map<int, std::string> id_to_token_;  // ID→token映射
    
    int unk_token_id_;
    int pad_token_id_;
    int eos_token_id_;
};
#endif //MSLITETOKENIZER_OPUS_MT_TOKENIZER_H
